import cv2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw

class ElementDetector:
    def __init__(self):
        pass
    
    def detect_edges(self, image_path, output_path=None):
        """
        使用边缘检测识别图片中的元素边界
        """
        # 读取图片
        img = cv2.imread(image_path)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # 边缘检测
        edges = cv2.Canny(gray, 50, 150)
        
        # 查找轮廓
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 绘制轮廓
        result_img = img.copy()
        
        detected_elements = []
        
        for i, contour in enumerate(contours):
            # 过滤小轮廓
            area = cv2.contourArea(contour)
            if area > 1000:  # 面积阈值
                # 获取边界框
                x, y, w, h = cv2.boundingRect(contour)
                
                # 绘制矩形框
                cv2.rectangle(result_img, (x, y), (x+w, y+h), (0, 255, 0), 3)
                
                # 添加标签
                cv2.putText(result_img, f'Element {i+1}', (x, y-10), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
                
                detected_elements.append({
                    'id': i+1,
                    'bbox': (x, y, w, h),
                    'area': area
                })
        
        # 显示结果
        plt.figure(figsize=(12, 8))
        plt.imshow(cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB))
        plt.axis('off')
        plt.title('元素边界检测结果')
        plt.show()
        
        if output_path:
            cv2.imwrite(output_path, result_img)
        
        return detected_elements

# 使用示例
# detector = ElementDetector()
# elements = detector.detect_edges("your_image.jpg", "elements_result.jpg")